#include <immintrin.h>
#include <pthread.h>
#include <cstring>
#include <iostream>
#include <random>
#include "mpi.h"
using namespace std;
#define N 4

inline bool in(int k, int begin, int end) {
    return k >= begin && k < end;
}

void createbuffer(float* buffer) {
    srand(20000214);
    for (int i = 0; i < N; i++) {
        for (int j = 0; j < N; j++) {
            buffer[i * N + j] = rand() % 100;
            // cout << buffer[i][j] << " ";
        }
        // cout << endl;
    }
    // cout << endl;
}

int main(int argc, char* argv[]) {
    __m512 t1, t2, t3, t4;
    int rank, size;
    double t11, t12;
    float *buffer, *row = new float[N];
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &size);
    int begin, end;
    MPI_Status status;
    begin = rank * (N / size);
    if (rank != size - 1) {
        end = begin + N / size;
    } else {
        end = N;
    }
    if (rank == 0) {
        buffer = new float[N * N];
        createbuffer(buffer);
        t11 = MPI_Wtime();
        for (int i = 1; i < size; i++) {
            if (i != size - 1) {
                MPI_Send(&buffer[i * N * (N / size)], N * (N / size), MPI_FLOAT,
                         i, i, MPI_COMM_WORLD);
            } else {
                MPI_Send(&buffer[i * N * (N / size)],
                         N * (N - (size - 1) * (N / size)), MPI_FLOAT, i, i,
                         MPI_COMM_WORLD);
            }
        }
    } else {
        int len = (end - begin) * N;
        buffer = new float[(end - begin) * N];
        MPI_Recv(buffer, len, MPI_FLOAT, 0, rank, MPI_COMM_WORLD, &status);
    }
    for (int k = 0; k < N; k++) {
        if (in(k, begin, end)) {
            float kk[16] = {
                buffer[(k - begin) * N + k], buffer[(k - begin) * N + k],
                buffer[(k - begin) * N + k], buffer[(k - begin) * N + k],
                buffer[(k - begin) * N + k], buffer[(k - begin) * N + k],
                buffer[(k - begin) * N + k], buffer[(k - begin) * N + k],
                buffer[(k - begin) * N + k], buffer[(k - begin) * N + k],
                buffer[(k - begin) * N + k], buffer[(k - begin) * N + k],
                buffer[(k - begin) * N + k], buffer[(k - begin) * N + k],
                buffer[(k - begin) * N + k], buffer[(k - begin) * N + k]};
            t1 = _mm512_loadu_ps(kk);
            for (int j = N - 16; j >= k; j -= 16) {
                t2 = _mm512_loadu_ps(&buffer[(k - begin) * N + j]);
                t3 = _mm512_div_ps(t2, t1);
                _mm512_storeu_ps(&buffer[(k - begin) * N + j], t3);
            }
            for (int j = k; j % 16 != N % 16; j++) {
                buffer[(k - begin) * N + j] /= kk[0];
            }
            memcpy(row, &buffer[(k - begin) * N], sizeof(float) * N);
        }
        int rk = k / (N / size);
        rk = rk >= size ? size - 1 : rk;
        MPI_Bcast(row, N, MPI_FLOAT, rk, MPI_COMM_WORLD);
        for (int i = begin; i < end; i++) {
            if (end <= k) {
                break;
            }
            if (i <= k) {
                continue;
            }
            float ik[16] = {
                buffer[(i - begin) * N + k], buffer[(i - begin) * N + k],
                buffer[(i - begin) * N + k], buffer[(i - begin) * N + k],
                buffer[(i - begin) * N + k], buffer[(i - begin) * N + k],
                buffer[(i - begin) * N + k], buffer[(i - begin) * N + k],
                buffer[(i - begin) * N + k], buffer[(i - begin) * N + k],
                buffer[(i - begin) * N + k], buffer[(i - begin) * N + k],
                buffer[(i - begin) * N + k], buffer[(i - begin) * N + k],
                buffer[(i - begin) * N + k], buffer[(i - begin) * N + k]};
            t1 = _mm512_loadu_ps(ik);
            for (int j = N - 16; j > k; j -= 16) {
                t2 = _mm512_loadu_ps(&buffer[(i - begin) * N + j]);
                t3 = _mm512_loadu_ps(&row[j]);
                t4 = _mm512_sub_ps(t2, _mm512_mul_ps(t1, t3));
                _mm512_storeu_ps(&buffer[(i - begin) * N + j], t4);
            }
            for (int j = k + 1; j % 16 != N % 16; j++) {
                buffer[(i - begin) * N + j] -=
                    buffer[(i - begin) * N + k] * row[j];
            }
            buffer[(i - begin) * N + k] = 0;
        }
        MPI_Barrier(MPI_COMM_WORLD);
    }
    if (rank == 0) {
        for (int i = 1; i < size; i++) {
            if (i != size - 1) {
                MPI_Recv(&buffer[N * i * (N / size)], N * (N / size), MPI_FLOAT,
                         i, i, MPI_COMM_WORLD, &status);
            } else {
                MPI_Recv(&buffer[N * i * (N / size)],
                         N * (N - (size - 1) * (N / size)), MPI_FLOAT, i, i,
                         MPI_COMM_WORLD, &status);
            }
        }
    } else {
        int len = (end - begin) * N;
        MPI_Send(buffer, len, MPI_FLOAT, 0, rank, MPI_COMM_WORLD);
    }
    if (rank == 0) {
        t12 = MPI_Wtime();
        cout << "row avx: " << (t12 - t11) * 1000 << "ms" << endl;
        // for (int i = 0, k = 0; i < N * N; i++) {
        //     cout << buffer[i] << ' ';
        //     if (++k % N == 0) {
        //         cout << endl;
        //     }
        // }
    }
    MPI_Finalize();
    delete[] buffer;
    delete[] row;
    return 0;
}